# 多轮规划和检查
# 改进prompt

import os
import numpy as np
import math
import ast
import json
import pygame
import imageio
from datetime import datetime
from openai import OpenAI
from apikey import api_key, base_url, model_name
import tempfile
import shutil
import sys
from PIL import Image, ImageDraw, ImageFont
import textwrap  


class LLMInterface:
    def __init__(self):
        self.api_key = api_key
        self.base_url = base_url
        self.model_name = model_name
        self.llm_round_count = 0  # 记录LLM调用轮数
        self.max_rounds = 5  # 最大轮数（一轮=一次规划+一次完成度判断）

        # 创建OpenAI客户端
        self.client = OpenAI(
            api_key=self.api_key,
            base_url=self.base_url,
        )

    def generate_plan(self, scene_desc, task_desc):
        """调用Qwen3模型生成任务规划"""
        if self.llm_round_count >= self.max_rounds:
            print(f"已达到最大LLM调用轮数({self.max_rounds}),不再生成新规划")
            return ""

        # 构建LLM提示
        prompt = self._build_llm_prompt(scene_desc, task_desc)

        print("-" * 20)
        print(prompt)
        print("-" * 20)

        # 调用Qwen3 API
        response = self._call_qwen3_api(prompt)

        print("-" * 20)
        print("完整响应:")
        print(response)
        print("-" * 20)

        return response

    def check_task_completion(self, scene_desc, task_desc):
        """调用LLM判断任务是否完成"""
        if self.llm_round_count >= self.max_rounds:
            print(f"已达到最大LLM调用轮数({self.max_rounds}),不再检查任务完成状态")
            return True

        # 构建提示
        prompt = (
            f"请根据当前场景和任务描述判断任务是否已完成。任务描述是:{task_desc}\n"
            f"当前场景描述如下:{scene_desc}\n"
            "请只回答'是'或'否',不要包含其他任何内容。"
        )

        print("-" * 20)
        print("任务完成检查提示:")
        print(prompt)
        print("-" * 20)

        # 调用API
        response = self._call_qwen3_api(prompt)

        print("-" * 20)
        print("任务完成检查响应:")
        print(response)
        print("-" * 20)

        # 解析响应
        return "是" in response or "yes" in response.lower()

    def _build_llm_prompt(self, scene_desc, task_desc):
        """构建LLM提示（y轴向上为正）"""
        # 添加关于物体尺寸和重力影响的说明
        physics_desc = (
            f"重要物理规则:\n"
            f"1. 桌子尺寸为0.5x0.5，箱子尺寸为0.2x0.2\n"
            f"2. 这是一个含重力的纵向平面，箱子必须放在桌子的上方或者地上，否则会因重力掉落\n"
            f"3. 箱子应该放在桌子或者地面上方一定距离的位置，而不是让机械臂移动到和桌子重合的坐标后松开\n"
            f"3. 抓取位置要和箱子重合"
            f"4. 判断箱子在桌子上的条件: \n"
            f"   - 箱子与桌子中心x坐标差小于(0.5 + 0.2)/2 = 0.35\n"
            f"   - 箱子y坐标-桌子y坐标= (0.5+0.5)/2=0.35\n"
            f"   - 误差范围小于0.05\n"
            f"5. 判断箱子在地上的条件: \n"
            f"   - 箱子y坐标等于其尺寸的一半\n"
            f"   - 误差范围小于0.05\n"
            f"6. 箱子之间会发生碰撞，所以注意放置的位置,也要注意要用一个把箱子提起来或者移动到上方后放下去的过程\n"
        )

        prompt = (
            "请将任务分解为机械臂的动作序列，并输出两部分内容：\n"
            "第一部分：分析任务情况（使用自然语言描述）\n"
            "   - 分析目标状态中各个对象的位置\n"
            "   - 分析需要移动的对象和移动顺序\n"
            "   - 分析哪些对象不需要移动\n"
            "   - 分析移动步骤中的中间坐标\n"
            "   - 考虑物理规则（箱子必须放在桌子上）\n"
            "第二部分：动作序列（在'-----'分隔符之后）\n"
            "   - 使用格式: [[动作类型, 目标x, 目标y, 目标旋转]]\n"
            "   - 动作类型只能是 'move', 'catch', 'release' 三种\n"
            "   - 'move': 机械臂移动到指定位置 (目标x, 目标y, 目标旋转)\n"
            "   - 'catch': 抓取对象,位姿参数无效\n"
            "   - 'release': 释放对象,位姿参数无效\n"
            "输出示例:\n"
            "分析部分...\n"
            "-----分隔符-----\n"
            '[["move", 1.5, 2.3, 0.78],\n'
            ' ["catch", 0, 0, 0],\n'
            ' ["move", 3.0, 1.0, 0.0],\n'
            ' ["release", 0, 0, 0]]\n'
            "规划要求:\n"
            "   - 规划结束时机械臂末端必须处于释放状态\n"
            "   - 规划时注意箱子必须放在桌子上或者地上\n"
            "   - 尽量减少不必要的操作\n"
            f"{physics_desc}\n"
            f"{scene_desc}\n"
            f"任务是:\n{task_desc}\n\n"
        )
        return prompt

    def _call_qwen3_api(self, prompt):
        """使用OpenAI风格的API调用Qwen3模型"""
        try:
            # 创建聊天完成请求
            completion = self.client.chat.completions.create(
                model=self.model_name,
                messages=[
                    {
                        "role": "system",
                        "content": "你是一个机械臂控制专家,负责将自然语言任务分解为机械臂动作序列。",
                    },
                    {"role": "user", "content": prompt},
                ],
                temperature=0.1,
                max_tokens=1000,
                # 对于Qwen3开源版模型,需要设置enable_thinking=False
                extra_body={"enable_thinking": False},
            )

            # 提取模型响应
            response = completion.choices[0].message.content
            return response

        except Exception as e:
            print(f"API调用异常: {e}")
            return ""


class ArmController:
    def __init__(self, env_state):
        self.env_state = env_state
        self.current_plan = []
        self.current_task_index = 0
        self.task_completed = False
        self.moving_to_target = False
        self.target_joint_angles = None
        self.current_action_description = "等待规划"  # 当前动作描述

    def set_plan(self, plan):
        """设置当前规划"""
        self.current_plan = plan
        self.current_task_index = 0
        self.task_completed = False
        self.moving_to_target = False
        self.target_joint_angles = None

        # 转换坐标系：将规划中的y坐标从"上正下负"转换为环境坐标系
        self._convert_plan_coordinates()

    def get_next_action(self):
        """获取下一个动作(关节角变化量)"""
        # 所有任务已完成
        if self.current_task_index >= len(self.current_plan):
            return None

        # 如果当前任务已完成,移动到下一个任务
        if self.task_completed:
            print("任务已完成,移动到下一个任务")
            self.current_task_index += 1
            self.task_completed = False
            self.moving_to_target = False
            self.target_joint_angles = None  # 重置目标关节角度

            # 检查是否还有任务
            if self.current_task_index >= len(self.current_plan):
                return None

        # 获取当前任务
        current_task = self.current_plan[self.current_task_index]
        action_type = current_task[0]

        # 处理移动任务
        if action_type == "move":
            # 如果是第一次进入移动任务,设置目标
            if not self.moving_to_target:
                # 设置目标位置和旋转
                self.current_target = {
                    "x": current_task[1],
                    "y": current_task[2],
                    "rotation": current_task[3],
                }
                self.moving_to_target = True
                self.current_action_description = (
                    f"移动至 ({current_task[1]:.2f}, {current_task[2]:.2f})"
                )

                # 计算逆向运动学
                self.target_joint_angles = self._inverse_kinematics(
                    self.current_target["x"],
                    self.current_target["y"],
                    self.current_target["rotation"],
                )
                print("逆向运动学结果:", self.target_joint_angles)

            # 检查是否到达目标位置
            if self._check_reached_target():
                self.task_completed = True
                return np.array([0, 0, 0, 0])  # 无动作
            else:
                # 计算关节角变化量
                delta_angles = self._calculate_delta_angles(self.target_joint_angles)
                # 返回关节角变化量
                return np.append(delta_angles, 0)  # 无夹爪动作

        # 处理抓取任务
        elif action_type == "catch":
            self.task_completed = True
            self.current_action_description = "抓取物体"
            return np.array([0, 0, 0, 1])  # 抓取动作

        # 处理释放任务
        elif action_type == "release":
            self.task_completed = True
            self.current_action_description = "释放物体"
            return np.array([0, 0, 0, 2])  # 释放动作

        return None

    def get_action_description(self):
        """获取当前动作描述"""
        return self.current_action_description

    def _convert_plan_coordinates(self):
        """将规划中的y坐标从'上正下负'转换为环境坐标系"""
        world_height = self.env_state["world_height"]
        for i, task in enumerate(self.current_plan):
            if task[0] == "move":
                # 转换y坐标: env_y = WORLD_HEIGHT - llm_y
                llm_y = task[2]
                env_y = world_height - llm_y
                self.current_plan[i][2] = env_y
                print(f"转换坐标: 任务{i+1} y从{llm_y}转换为{env_y}")

    def _inverse_kinematics(self, x, y, rotation):
        """
        逆向运动学:使用几何方法计算目标位姿对应的关节角度
        """
        print(f"正在进行移动到{(x,y,rotation)}的逆向运动学计算")
        # 基座位置
        base_x, base_y = self.env_state["base_pos"]

        # 连杆长度
        l1, l2, l3 = self.env_state["link_lengths"]

        # 计算末端执行器的实际位置(考虑基座偏移)
        target_x = x - base_x
        target_y = y - base_y

        # 计算第三个关节的位置(末端位置减去第三个连杆在目标方向上的分量)
        # 注意:第三个连杆的方向由目标旋转决定
        joint2_x = target_x - l3 * math.cos(rotation)
        joint2_y = target_y - l3 * math.sin(rotation)

        # 计算第二个关节到基座的距离
        distance_to_joint2 = math.sqrt(joint2_x**2 + joint2_y**2)

        # 检查是否可达
        if distance_to_joint2 > l1 + l2 or distance_to_joint2 < abs(l1 - l2):
            # 如果不可达,调整到最近可达点
            if distance_to_joint2 > l1 + l2:
                scale = (l1 + l2) / distance_to_joint2
            else:
                scale = abs(l1 - l2) / distance_to_joint2
            joint2_x *= scale
            joint2_y *= scale
            print(f"目标位置不可达,已调整到最近点: ({joint2_x}, {joint2_y})")

        # 计算第一个关节角度 q1
        # 从基座到第二个关节的向量角度
        alpha = math.atan2(joint2_y, joint2_x)

        # 使用余弦定理计算第二个关节角度 q2
        cos_q2 = (joint2_x**2 + joint2_y**2 - l1**2 - l2**2) / (2 * l1 * l2)
        cos_q2 = np.clip(cos_q2, -1.0, 1.0)  # 确保在有效范围内
        q2 = math.acos(cos_q2)

        # 计算第一个关节角度 q1
        # 使用两个可能的解(肘部向上和肘部向下)
        # 这里我们选择肘部向下的解(q2为负)
        beta = math.atan2(l2 * math.sin(q2), l1 + l2 * math.cos(q2))
        q1 = alpha - beta

        # 计算第三个关节角度 q3
        q3 = rotation - q1 - q2

        # 将角度限制在[-π, π]范围内
        q1 = self._wrap_angle(q1)
        q2 = self._wrap_angle(q2)
        q3 = self._wrap_angle(q3)

        print(f"逆向运动学计算完成: q1={q1:.4f}, q2={q2:.4f}, q3={q3:.4f}")
        return np.array([q1, q2, q3])

    def _calculate_delta_angles(self, target_angles):
        """
        计算关节角度变化量
        """
        delta = []
        current_angles = self.env_state["joint_angles"]
        for i in range(3):
            current = current_angles[i]
            target = target_angles[i]

            # 计算最小角度差
            diff = target - current
            diff = self._wrap_angle(diff)

            # 限制最大变化量
            max_step = 0.1  # 弧度
            if abs(diff) > max_step:
                diff = max_step if diff > 0 else -max_step

            delta.append(diff)

        return np.array(delta)

    def _check_reached_target(self):
        """检查机械臂是否到达目标位置"""
        # 获取当前末端位姿
        current_pos = self.env_state["end_effector_pos"]
        current_rotation = self.env_state["end_effector_rotation"]

        # 计算位置误差
        pos_error = np.linalg.norm(
            np.array(
                [
                    self.current_target["x"] - current_pos[0],
                    self.current_target["y"] - current_pos[1],
                ]
            )
        )

        # 计算旋转误差
        rot_error = abs(
            self._wrap_angle(self.current_target["rotation"] - current_rotation)
        )

        # 检查是否在允许的误差范围内
        position_threshold = 0.05  # 5厘米
        rotation_threshold = 0.087  # 约5度

        return pos_error < position_threshold and rot_error < rotation_threshold

    def _wrap_angle(self, angle):
        """
        将角度限制在[-π, π]范围内
        """
        while angle > np.pi:
            angle -= 2 * np.pi
        while angle < -np.pi:
            angle += 2 * np.pi
        return angle


def build_scene_state(env):
    """构建场景状态字典"""
    # 获取机械臂末端位姿并转换y坐标
    end_pos = env._forward_kinematics()
    rotation = sum(env.joint_angles)

    # 添加物体描述（转换y坐标）
    objects = []
    for table in env.tables:
        objects.append(
            {
                "type": "桌子",
                "x": table["pos"][0],
                "y": env.WORLD_HEIGHT - table["pos"][1],  # 转换y坐标
                "size": table["size"],
                "rotation": 0,
                "color": get_color_name(table["color"]),
            }
        )

    for box in env.boxes:
        objects.append(
            {
                "type": "箱子",
                "x": box["pos"][0],
                "y": env.WORLD_HEIGHT - box["pos"][1],  # 转换y坐标
                "size": box["size"],
                "rotation": 0,
                "color": get_color_name(box["color"]),
            }
        )

    # 构建场景状态字典
    scene_state = {
        "world_width": env.WORLD_WIDTH,
        "world_height": env.WORLD_HEIGHT,
        "base_pos": env.base_pos,
        "link_lengths": env.LINK_LENGTHS,
        "joint_angles": env.joint_angles.copy(),
        "end_effector_pos": end_pos,
        "end_effector_rotation": rotation,
        "objects": objects,
    }

    return scene_state


def build_scene_description(scene_state):
    """从场景状态构建场景描述文本（y轴向上为正）"""
    desc = f"现在的场景是:\n"
    desc += (
        f"环境尺寸为 ({scene_state['world_width']}, {scene_state['world_height']})\n"
    )

    # 机械臂末端位姿（已转换y坐标）
    end_pos = scene_state["end_effector_pos"]
    rotation = scene_state["end_effector_rotation"]
    desc += f"机械臂当前的末端位姿为 ({end_pos[0]:.2f}, {scene_state['world_height'] - end_pos[1]:.2f}, {rotation:.2f})\n"

    # 添加物体描述
    for obj in scene_state["objects"]:
        desc += (
            f"在位置 ({obj['x']:.2f}, {obj['y']:.2f}) "
            f"有一个{obj['type']}对象, 尺寸为 {obj['size'][0]}x{obj['size'][1]}, "
            f"颜色为 {obj['color']}。\n"
        )
    return desc


def parse_plan(plan_str):
    """解析规划字符串"""
    # 尝试提取JSON格式部分
    try:
        # 查找第一个[和最后一个]
        start_idx = plan_str.find("[")
        end_idx = plan_str.rfind("]")

        if start_idx != -1 and end_idx != -1 and end_idx > start_idx:
            json_str = plan_str[start_idx : end_idx + 1]
            print("尝试解析JSON:", json_str)
            return json.loads(json_str)
    except json.JSONDecodeError:
        pass

    # 尝试解析类似Python列表的字符串
    try:
        # 移除可能的多余字符
        clean_str = plan_str.replace("\n", "").replace(" ", "")
        # 尝试直接解析
        return ast.literal_eval(clean_str)
    except:
        # 尝试提取有效部分
        start = plan_str.find("[")
        end = plan_str.rfind("]") + 1
        if start != -1 and end != -1 and end > start:
            try:
                return ast.literal_eval(plan_str[start:end])
            except:
                pass

    print("无法解析规划:", plan_str)
    return []


def get_color_name(rgb):
    """
    将RGB颜色转换为名称
    """
    colors = {
        (255, 255, 255): "白色",
        (0, 0, 0): "黑色",
        (128, 128, 128): "灰色",
        (255, 0, 0): "红色",
        (0, 255, 0): "绿色",
        (0, 120, 255): "蓝色",
        (255, 255, 0): "黄色",
        (180, 0, 255): "紫色",
        (255, 100, 180): "粉色",
        (255, 150, 0): "橙色",
    }
    return colors.get(tuple(rgb), "未知颜色")


def add_text_to_image(
    image_array, text, position=(10, 10), font_size=20, color=(0, 0, 0), max_width=40
):
    """
    在图像上添加文字，支持自动换行
    max_width: 每行最大字符数
    """
    # 将numpy数组转换为PIL图像
    image = Image.fromarray(image_array)

    # 创建一个可绘制的对象
    draw = ImageDraw.Draw(image)

    try:
        # 尝试加载支持中文的字体
        # Windows 系统常用字体
        font_paths = [
            "C:/Windows/Fonts/simhei.ttf",  # 黑体
            "C:/Windows/Fonts/simsun.ttc",  # 宋体
            "C:/Windows/Fonts/msyh.ttc",  # 微软雅黑
            "C:/Windows/Fonts/msyhbd.ttc",  # 微软雅黑粗体
        ]

        # macOS 系统常用字体
        mac_font_paths = [
            "/System/Library/Fonts/PingFang.ttc",  # 苹方
            "/System/Library/Fonts/STHeiti Light.ttc",  # 华文黑体
            "/System/Library/Fonts/STHeiti Medium.ttc",  # 华文黑体
        ]

        # Linux 系统常用字体
        linux_font_paths = [
            "/usr/share/fonts/truetype/droid/DroidSansFallbackFull.ttf",  # Droid Sans
            "/usr/share/fonts/truetype/wqy/wqy-microhei.ttc",  # 文泉驿微米黑
        ]

        # 尝试加载所有可能的字体路径
        font = None
        for path in font_paths + mac_font_paths + linux_font_paths:
            try:
                font = ImageFont.truetype(path, font_size)
                break  # 找到可用字体就停止
            except IOError:
                continue

        # 如果没有找到任何字体,使用默认字体
        if font is None:
            font = ImageFont.load_default()
            print("警告:未找到中文字体,使用默认字体(可能不支持中文)")

    except Exception as e:
        print(f"字体加载错误: {e}")
        font = ImageFont.load_default()

    # 文本换行处理
    lines = []
    # 如果文本包含换行符，先按换行符分割
    paragraphs = text.split("\n")
    for para in paragraphs:
        # 对每个段落进行换行处理
        if para:
            # 使用textwrap进行换行
            wrapped_lines = textwrap.wrap(para, width=max_width)
            lines.extend(wrapped_lines)
        else:
            # 保留空行
            lines.append("")

    # 添加文字
    y_offset = position[1]
    for line in lines:
        draw.text((position[0], y_offset), line, fill=color, font=font)
        y_offset += font_size + 5  # 行间距

    # 将PIL图像转换回numpy数组
    return np.array(image)


# 测试代码
if __name__ == "__main__":
    from arm_env import ArmEnv

    # 创建环境
    env = ArmEnv()
    env.reset()

    # 创建LLM接口和机械臂控制器
    llm_interface = LLMInterface()

    # 创建临时文件夹保存帧
    temp_dir = tempfile.mkdtemp()
    frames = []
    frame_count = 0

    # 定义任务
    env.render()
    task_description = input("请输入任务描述: ")
    user_task = task_description  # 保存用户输入的指令

    # 初始化状态
    task_completed = False
    total_steps = 0
    max_total_steps = 20000  # 整个任务的最大步骤数

    # 主循环 - 支持多次规划
    while (
        not task_completed and llm_interface.llm_round_count < llm_interface.max_rounds
    ):
        # 增加LLM轮数计数
        llm_interface.llm_round_count += 1
        print(f"开始第 {llm_interface.llm_round_count} 轮规划")

        # 构建场景状态
        scene_state = build_scene_state(env)
        scene_description = build_scene_description(scene_state)

        # 创建控制器
        controller = ArmController(scene_state)

        # 生成规划
        response = llm_interface.generate_plan(scene_description, task_description)

        # 如果没有有效响应
        if not response:
            print("LLM未返回有效响应")
            break

        # 解析规划
        plan = parse_plan(response)
        print("生成的规划:")
        for i, step in enumerate(plan):
            print(f"步骤 {i+1}: {step}")

        # 如果没有生成有效规划
        if not plan:
            print("无法生成有效规划")
            break

        # 设置控制器规划
        controller.set_plan(plan)

        # 保存初始状态
        pygame.image.save(
            env.screen, os.path.join(temp_dir, f"frame_{frame_count:04d}.png")
        )
        frame_img = imageio.imread(
            os.path.join(temp_dir, f"frame_{frame_count:04d}.png")
        )

        # 添加文本信息
        text_info = (
            f"任务: {user_task}\n"
            f"轮数: {llm_interface.llm_round_count}/{llm_interface.max_rounds}\n"
            "状态: 规划开始"
        )
        frame_img = add_text_to_image(
            frame_img, text_info, position=(10, 10), font_size=20, color=(0, 0, 0)
        )
        frames.append(frame_img)
        frame_count += 1

        # 执行当前规划
        plan_running = True
        step_count = 0  # 当前规划内的步骤计数器
        max_steps_per_plan = 20000  # 每个规划的最大步骤限制

        while plan_running and step_count < max_steps_per_plan:
            step_count += 1
            total_steps += 1

            # 检查总步数限制
            if total_steps > max_total_steps:
                print(f"\n解算失败:超过最大总步骤限制({max_total_steps})")
                plan_running = False
                break

            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    plan_running = False
                    break

            # 更新场景状态
            scene_state = build_scene_state(env)
            controller.env_state = scene_state  # 更新控制器中的状态

            action = controller.get_next_action()
            # print(f"步骤 {step_count} - 获得动作:", action)

            # 没有更多动作,结束当前规划
            if action is None:
                print("当前规划执行完成")
                plan_running = False
                continue

            # 如果动作不为空,执行动作
            if np.any(action[:3] != 0) or action[3] != 0:
                env.step(action)
                env.render()

                # 保存当前帧并立即读取到内存
                pygame.image.save(
                    env.screen, os.path.join(temp_dir, f"frame_{frame_count:04d}.png")
                )
                frame_img = imageio.imread(
                    os.path.join(temp_dir, f"frame_{frame_count:04d}.png")
                )

                # 添加文本信息
                text_info = (
                    f"任务: {user_task}\n"
                    f"轮数: {llm_interface.llm_round_count}/{llm_interface.max_rounds}\n"
                    f"步骤: {step_count}\n"
                    f"动作: {controller.get_action_description()}"
                )
                frame_img = add_text_to_image(
                    frame_img,
                    text_info,
                    position=(10, 10),
                    font_size=20,
                    color=(0, 0, 0),
                )
                frames.append(frame_img)
                frame_count += 1
            else:
                # 空动作,复制前一帧
                frames.append(frames[-1])
                frame_count += 1

        # 检查任务完成状态
        scene_state = build_scene_state(env)
        scene_description = build_scene_description(scene_state)
        task_completed = llm_interface.check_task_completion(
            scene_description, task_description
        )
        print(f"任务完成状态: {task_completed}")
        if task_completed:
            print("任务完成")
            break

    print("任务终止。")
    # 如果循环结束但任务未完成
    if not task_completed:
        if llm_interface.llm_round_count >= llm_interface.max_rounds:
            print(f"\n解算失败:超过最大LLM调用轮数({llm_interface.max_rounds})")
            print("可能原因:")
            print("1. 任务过于复杂或表述不清晰")
            print("2. 环境状态变化导致规划失效")
            print("3. LLM无法生成有效解决方案")
        else:
            print("\n解算失败:未知原因")

    # 创建GIF
    gif_filename = f"simulation_{datetime.now().strftime('%Y%m%d_%H%M%S')}.gif"

    # 保存为GIF
    if frames:  # 确保有帧可保存
        imageio.mimsave(gif_filename, frames, fps=10, loop=0)
        print(f"仿真结果已保存为: {gif_filename}")
    else:
        print("没有生成仿真结果")

    # 清理临时文件
    shutil.rmtree(temp_dir)

    pygame.quit()

    # 如果任务未完成,退出程序
    if not task_completed:
        sys.exit(1)
